import cv2  
import numpy as np  
  
class TemplateMatcher:  
    def __init__(self, template_path, target_image_path):  
        self.template_path = template_path  
        self.target_image_path = target_image_path  
        self.templates = {}  
        self.chars = {}  
        self.target_img = None  
        self.resized_templates = {}  
  
    def load_templates(self):  
        char = {
            'minus': '-',
            'dot': '.',
            'zero': '0',  
            'one': '1',  
            'two': '2',  
            'three': '3',  
            'four': '4',  
            'five': '5',  
            'six': '6',  
            'seven': '7',  
            'eight': '8',  
            'nine': '9',  
        }
        for name in ['minus', 'dot', 'zero', 'one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine']:  # 可以扩展为从文件或配置中读取模板名称  
            template_path = f"{self.template_path}{name}.jpg"  
            self.templates[name] = cv2.imread(template_path, 0)  
            self.chars[name] = char[name]
  
    def load_target_image(self):
        if type(self.target_image_path) is str:
            self.target_img = cv2.imread(self.target_image_path, 0)
        else:
            self.target_img = self.target_image_path
  
    def resize_templates(self):  
        target_height, target_width = self.target_img.shape  
        for name, template in self.templates.items():  
            template_height, template_width = template.shape  
            scale_factor = target_height / template_height  
            resized_width = int(template_width * scale_factor)  
            resized_template = cv2.resize(template, (resized_width, target_height), interpolation=cv2.INTER_AREA)  
            self.resized_templates[name] = resized_template  
  
    def match_templates(self, threshold=0.7):  
        matches = {}

        # match all templates
        for name, template in self.resized_templates.items():  
            result = cv2.matchTemplate(self.target_img, template, cv2.TM_CCOEFF_NORMED)  
            loc = np.where(result >= threshold)  
            matches[name] = list(zip(loc[0], loc[1]))

        # match dot template
        dot_result = cv2.matchTemplate(self.target_img, self.templates['dot'], cv2.TM_CCOEFF_NORMED)  
        dot_loc = np.where(dot_result >= threshold)  
        matches['dot'] = list(zip(dot_loc[0], dot_loc[1]))
        return matches  
  
    def predict_characters(self, matches):  
        predicted_chars = []
        dot_predicted_chars = []
        target_height, target_width = self.target_img.shape  
        for name, positions in matches.items():
            if name != 'dot':
                for x, y in positions:  
                    if not any(abs(x - pos_x) < target_height // 4 and abs(y - pos_y) < target_width // 4 for pos_x, pos_y in [pos for char, pos in predicted_chars]):  
                        predicted_chars.append((self.chars[name], (x, y)))  
            else:
                for x, y in positions:  
                    if not any(abs(x - pos_x) < target_height // 4 and abs(y - pos_y) < target_width // 4 for pos_x, pos_y in [pos for char, pos in dot_predicted_chars]):  
                        dot_predicted_chars.append((self.chars[name], (x, y)))
        predicted_chars = predicted_chars + dot_predicted_chars
        predicted_chars.sort(key=lambda x: x[1][1])  
        return predicted_chars  

    def get_predicted_string(self, predicted_chars):  
        return ''.join(char for char, pos in predicted_chars)  
    
  
    def process(self, threshold=0.7):  
        self.load_templates()  
        self.load_target_image()  
        self.resize_templates()  
        matches = self.match_templates(threshold)  
        predicted_chars = self.predict_characters(matches)  
        predicted_string = self.get_predicted_string(predicted_chars)  
        # print(f"Predicted characters: {predicted_string}")  
        return predicted_string  
  
# 使用示例  
# template_matcher = TemplateMatcher('image/template/', 'image/raw_image/HSV-2-Digital-4.jpg')  
# predicted_string = template_matcher.process(0.7)